#!/usr/bin/evn python
# coding:utf-8
import os
import numpy as np

try:
    import xml.etree.cElementTree as ET
except ImportError:
    import xml.etree.ElementTree as ET
import sys

train_name = "./UW20170906/ImageSets/Main/trainval.txt"
test_name = "./UW20170906/ImageSets/Main/test.txt"
file_train = open(train_name, 'w')
file_train.truncate()
file_train.close()
file_train = open(train_name, 'a')

file_test = open(test_name, 'w')
file_test.truncate()
file_test.close()
file_test = open(test_name, 'a')

train_test_ratio = 0.0

for root, dirs, files in os.walk("./UW20170906/Annotations/", topdown=False):
    print(files)
for file in files:
    print(os.path.join(root, file))
    f = os.path.join(root, file)
    tree = ET.parse(f)  # 打开xml文档wi
    root_node = tree.getroot()  # 获得root节点
    filename = root_node.find('filename').text
    rand = np.random.random()
    if rand > train_test_ratio:
        file_train.write(filename+"\n")
    else:
        file_test.write(filename+"\n")
file_test.close()
file_train.close()
